from typing import Any, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, DatasetWriter, pause, resume, skip

AnySample = TypeVar('AnySample')


class DelegatingIterator(DatasetIterator[AnySample]):
    def __init__(self, delegated_iterator: DatasetIterator[AnySample]):
        self._delegated_iterator = delegated_iterator

    def __next__(self) -> AnySample:
        return next(self._delegated_iterator)

    def __skip__(self, samples_n: int):
        skip(self._delegated_iterator, samples_n)

    def __pause__(self) -> Any:
        return pause(self._delegated_iterator)


class DelegatingDataset(Dataset[AnySample]):
    def __init__(self, delegated_dataset: Dataset[AnySample]):
        self._delegated_dataset = delegated_dataset

    def __len__(self) -> int:
        return len(self._delegated_dataset)

    def __getitem__(self, index: int, /) -> AnySample:
        return self._delegated_dataset[index]

    def __iter__(self) -> DatasetIterator[AnySample]:
        return iter(self._delegated_dataset)

    def __resume__(self, state: Any | None = None, /) -> DatasetIterator[AnySample]:
        return resume(self._delegated_dataset, state)


class DelegatingDatasetWriter(DatasetWriter[AnySample]):
    def __init__(self, delegated_dataset_writer: DatasetWriter[AnySample]):
        self._delegated_dataset_writer = delegated_dataset_writer

    def write(self, sample: AnySample, /):
        self._delegated_dataset_writer.write(sample)

    def flush(self):
        self._delegated_dataset_writer.flush()

    def close(self):
        self._delegated_dataset_writer.close()
